#include "BlobTracker.h"

double cosAngle( CvPoint* pt1, CvPoint* pt2, CvPoint* pt0 )
{
    double dx1 = pt1->x - pt0->x;
    double dy1 = pt1->y - pt0->y;
    double dx2 = pt2->x - pt0->x;
    double dy2 = pt2->y - pt0->y;
    return (dx1*dx2 + dy1*dy2)/sqrt((dx1*dx1 + dy1*dy1)*(dx2*dx2 + dy2*dy2) + 1e-10);
}

unsigned int getTag(CvSeq *curCont, int level=0)
{

	unsigned int num = 0;
	unsigned int sum = 0;

	if(level > 5)
		return 0;

	for( ; curCont != 0; curCont = curCont->h_next )	
	{
		num ++;

		if(curCont->v_next)
			sum += getTag(curCont->v_next, level+1);

		if(num > 9)
			return 9;

	}

	sum += num * (unsigned int)pow(10.0, level);

	return sum;
}

Blob::Blob()
{
	area = 0.0f;
	weight = 0.0f;
}

Blob::~Blob()
{
}


BlobTracker::BlobTracker(void)
{
	currentID = 1;
	extraIDs = 0;
	ghost_frames = 0;
	reject_distance = 250;
	reject_min = 20;
	reject_max = 80;
	min_displacement = 0.5f;
}

void BlobTracker::find_blob(IplImage* src)
{

	blob_list.clear();
	CvMemStorage* storage = cvCreateMemStorage(0);
	CvSeq* cont = 0; 
	CvBox2D32f box;
	float halfx,halfy;
	Blob blob;
	CvSeq *result;
    //CvSeq *squares = cvCreateSeq( 0, sizeof(CvSeq), sizeof(CvPoint), storage );

	double s, t;
	unsigned int i;

	bool isSquare = false;
	IplImage *img = cvCloneImage(src);

	cvFindContours( img, storage, &cont, sizeof(CvContour), CV_RETR_TREE, CV_CHAIN_APPROX_SIMPLE );

	for( ; cont != 0; cont = cont->h_next )	{
		int count = cont->total; // This is number point in contour

		// First we check to see if this contour looks like a square.. 
		isSquare = false;
        result = cvApproxPoly( cont, sizeof(CvContour), storage,
            CV_POLY_APPROX_DP, cvContourPerimeter(cont)*0.02, 0 );

        if( result->total == 4 &&
            fabs(cvContourArea(result,CV_WHOLE_SEQ)) > 1000 &&
            cvCheckContourConvexity(result) )
        {
			s = 0;
            
            for( i = 0; i < 5; i++ )
            {
                // find minimum angle between joint
                // edges (maximum of cosine)
                if( i >= 2 )
                {
                    t = fabs(cosAngle(
                    (CvPoint*)cvGetSeqElem( result, i ),
                    (CvPoint*)cvGetSeqElem( result, i-2 ),
                    (CvPoint*)cvGetSeqElem( result, i-1 )));
                    s = s > t ? s : t;
                }
            }
            
            // if cosines of all angles are small
            // (all angles are ~90 degree) then this is a square.. 
            if( s < 0.5 )
			{
	

				box = cvMinAreaRect2(cont, storage);		
				blob.center.X = box.center.x;
				blob.center.Y = box.center.y;
				blob.angle = box.angle;

				halfx = box.size.width*0.5f;
				halfy = box.size.height*0.5f;
				blob.box.upperLeftCorner.set(box.center.x-halfx,box.center.y-halfy);
				blob.box.lowerRightCorner.set(box.center.x+halfx,box.center.y+halfy);

				blob.area = blob.box.getArea();

				// FIXME: it might be nice if we could get the actual weight.. 
				// It also might be nice to find the weighted center..

				blob.weight = 0;


				// use v_next.. 

				if(cont->v_next)
					blob.tagID = getTag(cont->v_next);

				printf("Square Detected %d\n", blob.tagID);

				blob_list.push_back(blob);

				isSquare = true;
			}
		}

		// fallback, if it's a regular blob.
		if(!isSquare)
		{
			// Number point must be more than or equal to 6 (for cvFitEllipse_32f).    
			if( count >= 6)
			{
				// Fits ellipse to current contour.
				box = cvFitEllipse2(cont);	
			} else {
				box = cvMinAreaRect2(cont, storage);		
			}
			blob.center.X = box.center.x;
			blob.center.Y = box.center.y;
			blob.angle = box.angle;

			halfx = box.size.width*0.5f;
			halfy = box.size.height*0.5f;
			blob.box.upperLeftCorner.set(box.center.x-halfx,box.center.y-halfy);
			blob.box.lowerRightCorner.set(box.center.x+halfx,box.center.y+halfy);

			blob.area = blob.box.getArea();
			
			// FIXME: it might be nice if we could get the actual weight.. 
			// It also might be nice to
			blob.weight = 0;
			blob.tagID = 0;

			int h=(int)blob.box.getHeight(), w=(int)blob.box.getWidth();
			if(w >= reject_min && h >= reject_min && w < reject_max && h < reject_max)
				blob_list.push_back(blob);
		}
	}		// end cont for loop
	cvReleaseImage(&img);
	cvReleaseMemStorage(&storage);

}

void BlobTracker::track_blob(void)
{
	size_t i, j, k;

	history.push_back(current);

	if(history.size() > HISTORY_FRAMES) 
	{
		history.erase(history.begin());
	}

	current.clear();

	size_t numblobs = blob_list.size();

	for(i=0; i<numblobs; i++)
	{
		current.push_back(Finger(blob_list[i]));
	}

	std::vector<Finger> *prev = &history[history.size()-1];

	size_t cursize = current.size();
	size_t prevsize = (*prev).size();

	for(i=0; i<cursize; i++)
	{
		current[i].error.clear();
		current[i].closest.clear();

		for(j=0; j<prevsize; j++)
		{
			float error = 0.0f;
			error = getError((*prev)[j], current[i]);
			current[i].error.push_back(error);
			current[i].closest.push_back(j);
		}
	}

	// sort so we can make a list of the closest blobs in the previous frame..
	for(i=0; i<cursize; i++)
	{
		// Bubble sort closest.
		for(j=0; j<prevsize; j++)
		{
			for(k=0; k<prevsize-1-j; k++)
			{
				// ugly as hell, I know.
				
				if(current[i].error[current[i].closest[k+1]] < current[i].error[current[i].closest[k]]) 
				{
				  int tmp = current[i].closest[k];			// swap
				  current[i].closest[k] = current[i].closest[k+1];
				  current[i].closest[k+1] = tmp;
				}
			}
		}
	}

	// Generate a matrix of all the possible choices

	ids.clear();

	for(i=0; i<cursize; i++)
	{
		ids.push_back(-1);
	}

	extraIDs = cursize - prevsize;
	if(extraIDs < 0)
		extraIDs = 0;
	matrix.clear();

	if(cursize <= 4)
		numcheck = 4;
	else if(cursize <= 6)
		numcheck = 3;
	else if(cursize <= 10)
		numcheck = 2;
	else
		numcheck = 1;

	if(prevsize < numcheck)
		numcheck = prevsize;
	
	if(current.size() > 0)
		cal_matrix(0);


	unsigned int num_results = matrix.size();

	// loop through all the potential ID configurations and find one with lowest error
	float best_error = 99999, error;
	int best_error_ndx = -1;

	for(j=0; j<num_results; j++)
	{
		error = 0;
		// get the error for each blob and sum
		for(i=0; i<cursize; i++)			
		{
			Finger *f = 0;

			if(matrix[j][i] != -1) 
			{
				error += current[i].error[matrix[j][i]];
			}
		}

		if(error < best_error)
		{
			best_error = error;
			best_error_ndx = j;		
		}
	}

	if(best_error_ndx != -1)
	{
		for(i=0; i<cursize; i++)
		{
			if(matrix[best_error_ndx][i] != -1){
				current[i].ID = (*prev)[matrix[best_error_ndx][i]].ID;
			}else{
				current[i].ID = -1;
			}

			if(current[i].ID != -1)
			{
				Finger *oldfinger = &(*prev)[matrix[best_error_ndx][i]];
				current[i].delta = (current[i].center - oldfinger->center);
				current[i].deltaArea = current[i].area - oldfinger->area;
				current[i].predictedPos = current[i].center + current[i].delta;
				current[i].displacement = oldfinger->displacement + current[i].delta;
				current[i].last_center = oldfinger->center;
			} else {
				current[i].delta = vector2df(0, 0);
				current[i].deltaArea = 0;
				current[i].predictedPos = current[i].center;
				current[i].displacement = vector2df(0.0f, 0.0f);
				current[i].last_center = vector2df(-1,-1);
			}
		}
		//printf("Best index = %d\n", best_error_ndx);
	}

}

void BlobTracker::cal_matrix(int start)
{
	if (start == ids.size()) 
  {

		  matrix.push_back(ids);

  }
  else 
  {
	  size_t numchecked=0;

	  for(size_t i=0; i<current[start].closest.size(); i++)
	  {
		if(current[start].error[current[start].closest[i]] > reject_distance)
			break;

		ids[start] = current[start].closest[i];
		if(check_valid(start))
		{
			cal_matrix(start+1);
			numchecked++;

		}

		if(numchecked >= numcheck)
			break;
	  }

	  if(extraIDs > 0)
	  {
			ids[start] = -1;		// new ID
			if(check_valid_new(start))
			{
				cal_matrix(start+1);
			}
	  }
  }
}

bool BlobTracker::check_valid(int start)
{
	for(int i=0; i<start; i++)
	{
	  // check to see whether this ID exists already
	  if(ids[i] == ids[start])
		  return false;
	}

	return true;
}

bool BlobTracker::check_valid_new(int start)
{
	int newidcount = 1;

	for(int i=0; i<start; i++)
	{
	  if(ids[i] == -1)
		  newidcount ++;
	}

	if(newidcount > extraIDs)		//extraIDs > 0 
	  return false;

	return true;
}

float BlobTracker::getError(Finger &old, Finger &cur)
{
	vector2df dev = cur.center - old.center;

	return (float) dev.getLength();

}

void BlobTracker::gather_events(void)
{
	std::vector<Finger> *prev = &history[history.size()-1];

	const int cursize = (unsigned int) current.size();
	const int prevsize = (unsigned int) (*prev).size();
	int i, j;

	// assign ID's for any blobs that are new this frame (ones that didn't get 
	// matched up with a blob from the previous frame).
	for(i=0; i<cursize; i++)
	{
		if(current[i].ID == -1)	
		{
			current[i].ID = currentID;

			currentID ++;
			if(currentID >= 65535)
				currentID = 0;

			do_touch(current[i].getTouchData());
		} else {
			if (current[i].displacement.getLength() >= min_displacement) {
				do_update(current[i].getTouchData());
				current[i].displacement = vector2df(0.0f, 0.0f);
			}
		}
	}

	// if a blob disappeared this frame, send a finger up event
	for(i=0; i<prevsize; i++)		// for each one in the last frame, see if it still exists in the new frame.
	{
		bool found = false;
		for(j=0; j<cursize; j++)
		{
			if(current[j].ID == (*prev)[i].ID)
			{
				found = true;
				break;
			}
		}

		if(!found)
		{
			

			if(ghost_frames == 0)
			{
				do_untouch((*prev)[i].getTouchData());
			} else if((*prev)[i].markedForDeletion)
			{
				(*prev)[i].framesLeft -= 1;
				if((*prev)[i].framesLeft <= 0)
					do_untouch((*prev)[i].getTouchData());
				else
					current.push_back((*prev)[i]);	// keep it around until framesleft = 0
			} else
			{
				(*prev)[i].markedForDeletion = true;
				(*prev)[i].framesLeft = ghost_frames;
				current.push_back((*prev)[i]);	// keep it around until framesleft = 0
			}

		}
	}
}

void BlobTracker::do_touch(TouchData tdata)
{
	tdata.tEvent = TOUCH_EVENT;
	vec_event.push_back(tdata);
}

void BlobTracker::do_untouch(TouchData tdata)
{
	tdata.tEvent = UNTOUCH_EVENT;
	vec_event.push_back(tdata);
}

void BlobTracker::do_update(TouchData tdata)
{
	tdata.tEvent = UPDATE_EVENT;
	vec_event.push_back(tdata);
}

std::vector<TouchData> BlobTracker::process(IplImage* img)
{
	vec_event.clear();
	find_blob(img);
	track_blob();
	gather_events();
	return vec_event;

}

BlobTracker::~BlobTracker(void)
{
}